{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9OEXtK6c6lyz"
      },
      "source": [
        "# Movinet Streaming Model Tutorial\n",
        "\n",
        "This tutorial trains a [MoViNet](https://arxiv.org/abs/2103.11511) with `A0` configuration from the TensorFlow Model Garden package (tensorflow-models).\n",
        "\n",
        "[Model Garden](https://www.tensorflow.org/tfmodels) contains a collection of state-of-the-art models, implemented with TensorFlow's high-level APIs. The implementations demonstrate the best practices for modeling, letting users to take full advantage of TensorFlow for their research and product development.\n",
        "\n",
        "**Streaming Models**:\n",
        "\n",
        "Streaming models implement causal (2+1)D convolutions with stream buffers. Streaming models use (2+1)D convolution instead of 3D to utilize optimized tf.nn.conv2d operations, which offer fast inference on CPU. Streaming models can be run on individual frames or on larger video clips like base models.\n",
        "\n",
        "**Note**: A3, A4, and A5 models use a positional encoding in the squeeze-excitation blocks, while A0, A1, and A2 do not. For the smaller models, accuracy is unaffected without positional encoding, while for the larger models accuracy is significantly worse without positional encoding.\n",
        "\n",
        "**Dataset:** [UCF_101](https://www.tensorflow.org/datasets/catalog/ucf101)\n",
        "* A 101-label video classification dataset\n",
        "\n",
        "**This tutorial demonstrates how to:**\n",
        "\n",
        "* Use models from the TensorFlow Models package.\n",
        "* Train/Fine-tune a pre-built [MoViNet](https://arxiv.org/abs/2103.11511) for Video Classification.\n",
        "* Export the trained/tuned MoViNet-A0-Stream model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "560IfSkP7fld"
      },
      "source": [
        "## Install cecessary libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MVX4IofucuZK"
      },
      "outputs": [],
      "source": [
        "!pip install -U -q \"tf-models-official\"\n",
        "\n",
        "\n",
        "# Install the mediapy package for visualizing images/videos.\n",
        "# See https://github.com/google/mediapy\n",
        "!command -v ffmpeg \u003e/dev/null || (apt update \u0026\u0026 apt install -y ffmpeg)\n",
        "!pip install -q mediapy remotezip\n",
        "!pip install -U -q git+https://github.com/tensorflow/docs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K-clGdjV7nF5"
      },
      "source": [
        "## Import necessary libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AklLemGhc_Jm"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import tqdm\n",
        "import random\n",
        "import pathlib\n",
        "import imageio\n",
        "import itertools\n",
        "import collections\n",
        "\n",
        "import cv2\n",
        "import numpy as np\n",
        "import remotezip as rz\n",
        "import seaborn as sns\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "from tensorflow_docs.vis import embed\n",
        "\n",
        "import keras\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "from tensorflow.keras import layers\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.losses import SparseCategoricalCrossentropy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "epVTJ5BH7ps3"
      },
      "source": [
        "## Import the MoViNet model from TensorFlow Models (tf-models-official)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xli9cfvw7s3t"
      },
      "outputs": [],
      "source": [
        "from official.projects.movinet.modeling import movinet\n",
        "from official.projects.movinet.modeling import movinet_model\n",
        "from official.projects.movinet.tools import export_saved_model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pbMYvcJX76KP"
      },
      "source": [
        "## Download Subdataset of UCF_101"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "F5XwQuLzd5Id"
      },
      "outputs": [],
      "source": [
        "# @title Helper functions for loading data and visualizing\n",
        "def list_files_per_class(zip_url):\n",
        "  \"\"\"\n",
        "    List the files in each class of the dataset given the zip URL.\n",
        "\n",
        "    Args:\n",
        "      zip_url: URL from which the files can be unzipped.\n",
        "\n",
        "    Return:\n",
        "      files: List of files in each of the classes.\n",
        "  \"\"\"\n",
        "  files = []\n",
        "  with rz.RemoteZip(URL) as zip:\n",
        "    for zip_info in zip.infolist():\n",
        "      files.append(zip_info.filename)\n",
        "  return files\n",
        "\n",
        "def get_class(fname):\n",
        "  \"\"\"\n",
        "    Retrieve the name of the class given a filename.\n",
        "\n",
        "    Args:\n",
        "      fname: Name of the file in the UCF101 dataset.\n",
        "\n",
        "    Return:\n",
        "      Class that the file belongs to.\n",
        "  \"\"\"\n",
        "  return fname.split('_')[-3]\n",
        "\n",
        "def get_files_per_class(files):\n",
        "  \"\"\"\n",
        "    Retrieve the files that belong to each class.\n",
        "\n",
        "    Args:\n",
        "      files: List of files in the dataset.\n",
        "\n",
        "    Return:\n",
        "      Dictionary of class names (key) and files (values).\n",
        "  \"\"\"\n",
        "  files_for_class = collections.defaultdict(list)\n",
        "  for fname in files:\n",
        "    class_name = get_class(fname)\n",
        "    files_for_class[class_name].append(fname)\n",
        "  return files_for_class\n",
        "\n",
        "def download_from_zip(zip_url, to_dir, file_names):\n",
        "  \"\"\"\n",
        "    Download the contents of the zip file from the zip URL.\n",
        "\n",
        "    Args:\n",
        "      zip_url: Zip URL containing data.\n",
        "      to_dir: Directory to download data to.\n",
        "      file_names: Names of files to download.\n",
        "  \"\"\"\n",
        "  with rz.RemoteZip(zip_url) as zip:\n",
        "    for fn in tqdm.tqdm(file_names):\n",
        "      class_name = get_class(fn)\n",
        "      zip.extract(fn, str(to_dir / class_name))\n",
        "      unzipped_file = to_dir / class_name / fn\n",
        "\n",
        "      fn = pathlib.Path(fn).parts[-1]\n",
        "      output_file = to_dir / class_name / fn\n",
        "      unzipped_file.rename(output_file,)\n",
        "\n",
        "def split_class_lists(files_for_class, count):\n",
        "  \"\"\"\n",
        "    Returns the list of files belonging to a subset of data as well as the remainder of\n",
        "    files that need to be downloaded.\n",
        "\n",
        "    Args:\n",
        "      files_for_class: Files belonging to a particular class of data.\n",
        "      count: Number of files to download.\n",
        "\n",
        "    Return:\n",
        "      split_files: Files belonging to the subset of data.\n",
        "      remainder: Dictionary of the remainder of files that need to be downloaded.\n",
        "  \"\"\"\n",
        "  split_files = []\n",
        "  remainder = {}\n",
        "  for cls in files_for_class:\n",
        "    split_files.extend(files_for_class[cls][:count])\n",
        "    remainder[cls] = files_for_class[cls][count:]\n",
        "  return split_files, remainder\n",
        "\n",
        "def download_ufc_101_subset(zip_url, num_classes, splits, download_dir):\n",
        "  \"\"\"\n",
        "    Download a subset of the UFC101 dataset and split them into various parts, such as\n",
        "    training, validation, and test.\n",
        "\n",
        "    Args:\n",
        "      zip_url: Zip URL containing data.\n",
        "      num_classes: Number of labels.\n",
        "      splits: Dictionary specifying the training, validation, test, etc. (key) division of data\n",
        "              (value is number of files per split).\n",
        "      download_dir: Directory to download data to.\n",
        "\n",
        "    Return:\n",
        "      dir: Posix path of the resulting directories containing the splits of data.\n",
        "  \"\"\"\n",
        "  files = list_files_per_class(zip_url)\n",
        "  for f in files:\n",
        "    tokens = f.split('/')\n",
        "    if len(tokens) \u003c= 2:\n",
        "      files.remove(f) # Remove that item from the list if it does not have a filename\n",
        "\n",
        "  files_for_class = get_files_per_class(files)\n",
        "\n",
        "  classes = list(files_for_class.keys())[:num_classes]\n",
        "\n",
        "  for cls in classes:\n",
        "    new_files_for_class = files_for_class[cls]\n",
        "    random.shuffle(new_files_for_class)\n",
        "    files_for_class[cls] = new_files_for_class\n",
        "\n",
        "  # Only use the number of classes you want in the dictionary\n",
        "  files_for_class = {x: files_for_class[x] for x in list(files_for_class)[:num_classes]}\n",
        "\n",
        "  dirs = {}\n",
        "  for split_name, split_count in splits.items():\n",
        "    print(split_name, \":\")\n",
        "    split_dir = download_dir / split_name\n",
        "    split_files, files_for_class = split_class_lists(files_for_class, split_count)\n",
        "    download_from_zip(zip_url, split_dir, split_files)\n",
        "    dirs[split_name] = split_dir\n",
        "\n",
        "  return dirs\n",
        "\n",
        "def format_frames(frame, output_size):\n",
        "  \"\"\"\n",
        "    Pad and resize an image from a video.\n",
        "\n",
        "    Args:\n",
        "      frame: Image that needs to resized and padded.\n",
        "      output_size: Pixel size of the output frame image.\n",
        "\n",
        "    Return:\n",
        "      Formatted frame with padding of specified output size.\n",
        "  \"\"\"\n",
        "  frame = tf.image.convert_image_dtype(frame, tf.float32)\n",
        "  frame = tf.image.resize_with_pad(frame, *output_size)\n",
        "  return frame\n",
        "\n",
        "def frames_from_video_file(video_path, n_frames, output_size = (172,172), frame_step = 15):\n",
        "  \"\"\"\n",
        "    Creates frames from each video file present for each category.\n",
        "\n",
        "    Args:\n",
        "      video_path: File path to the video.\n",
        "      n_frames: Number of frames to be created per video file.\n",
        "      output_size: Pixel size of the output frame image.\n",
        "\n",
        "    Return:\n",
        "      An NumPy array of frames in the shape of (n_frames, height, width, channels).\n",
        "  \"\"\"\n",
        "  # Read each video frame by frame\n",
        "  result = []\n",
        "  src = cv2.VideoCapture(str(video_path))\n",
        "\n",
        "  video_length = src.get(cv2.CAP_PROP_FRAME_COUNT)\n",
        "\n",
        "  need_length = 1 + (n_frames - 1) * frame_step\n",
        "\n",
        "  if need_length \u003e video_length:\n",
        "    start = 0\n",
        "  else:\n",
        "    max_start = video_length - need_length\n",
        "    start = random.randint(0, max_start + 1)\n",
        "\n",
        "  src.set(cv2.CAP_PROP_POS_FRAMES, start)\n",
        "  # ret is a boolean indicating whether read was successful, frame is the image itself\n",
        "  ret, frame = src.read()\n",
        "  result.append(format_frames(frame, output_size))\n",
        "\n",
        "  for _ in range(n_frames - 1):\n",
        "    for _ in range(frame_step):\n",
        "      ret, frame = src.read()\n",
        "    if ret:\n",
        "      frame = format_frames(frame, output_size)\n",
        "      result.append(frame)\n",
        "    else:\n",
        "      result.append(np.zeros_like(result[0]))\n",
        "  src.release()\n",
        "  result = np.array(result)[..., [2, 1, 0]]\n",
        "\n",
        "  return result\n",
        "\n",
        "def to_gif(images):\n",
        "  converted_images = np.clip(images * 255, 0, 255).astype(np.uint8)\n",
        "  imageio.mimsave('./animation.gif', converted_images, fps=10)\n",
        "  return embed.embed_file('./animation.gif')\n",
        "\n",
        "\n",
        "class FrameGenerator:\n",
        "  def __init__(self, path, n_frames, training = False):\n",
        "    \"\"\" Returns a set of frames with their associated label.\n",
        "\n",
        "      Args:\n",
        "        path: Video file paths.\n",
        "        n_frames: Number of frames.\n",
        "        training: Boolean to determine if training dataset is being created.\n",
        "    \"\"\"\n",
        "    self.path = path\n",
        "    self.n_frames = n_frames\n",
        "    self.training = training\n",
        "    self.class_names = sorted(set(p.name for p in self.path.iterdir() if p.is_dir()))\n",
        "    self.class_ids_for_name = dict((name, idx) for idx, name in enumerate(self.class_names))\n",
        "\n",
        "  def get_files_and_class_names(self):\n",
        "    video_paths = list(self.path.glob('*/*.avi'))\n",
        "    classes = [p.parent.name for p in video_paths]\n",
        "    return video_paths, classes\n",
        "\n",
        "  def __call__(self):\n",
        "    video_paths, classes = self.get_files_and_class_names()\n",
        "\n",
        "    pairs = list(zip(video_paths, classes))\n",
        "\n",
        "    if self.training:\n",
        "      random.shuffle(pairs)\n",
        "\n",
        "    for path, name in pairs:\n",
        "      video_frames = frames_from_video_file(path, self.n_frames)\n",
        "      label = self.class_ids_for_name[name] # Encode labels\n",
        "      yield video_frames, label"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JbHsejeveDFW"
      },
      "outputs": [],
      "source": [
        "# Helper functions below used are taken from following tutorials\n",
        "# https://www.tensorflow.org/tutorials/video/video_classification\n",
        "# https://www.tensorflow.org/tutorials/video/transfer_learning_with_movinet\n",
        "\n",
        "URL = 'https://storage.googleapis.com/thumos14_files/UCF101_videos.zip'\n",
        "download_dir = pathlib.Path('./UCF101_subset/')\n",
        "subset_paths = download_ufc_101_subset(URL,\n",
        "                        num_classes = 10,\n",
        "                        splits = {\"train\": 40, \"val\": 10, \"test\": 10},\n",
        "                        download_dir = download_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Riw5ZQDx8aA0"
      },
      "source": [
        "## Prepare train, valid and test dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CJ2xmsdVeGMM"
      },
      "outputs": [],
      "source": [
        "batch_size = 10\n",
        "num_frames = 8\n",
        "\n",
        "CLASSES = sorted(os.listdir('./UCF101_subset/train'))\n",
        "\n",
        "output_signature = (tf.TensorSpec(shape = (None, None, None, 3), dtype = tf.float32),\n",
        "                    tf.TensorSpec(shape = (), dtype = tf.int16))\n",
        "\n",
        "train_ds = tf.data.Dataset.from_generator(FrameGenerator(subset_paths['train'], num_frames, training = True),\n",
        "                                          output_signature = output_signature)\n",
        "train_ds = train_ds.batch(batch_size)\n",
        "\n",
        "val_ds = tf.data.Dataset.from_generator(FrameGenerator(subset_paths['val'], num_frames),\n",
        "                                          output_signature = output_signature)\n",
        "val_ds = val_ds.batch(batch_size)\n",
        "\n",
        "test_ds = tf.data.Dataset.from_generator(FrameGenerator(subset_paths['test'], num_frames),\n",
        "                                         output_signature = output_signature)\n",
        "test_ds = test_ds.batch(batch_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HuyVEA5l8myh"
      },
      "source": [
        "### Check the prepared training batch of the data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5tPCI35TeihA"
      },
      "outputs": [],
      "source": [
        "for frames, labels in train_ds.take(1):\n",
        "  print(f\"Shape: {frames.shape}\")\n",
        "  print(f\"Label: {labels.shape}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aNk_b5gk8s-g"
      },
      "source": [
        "## Build the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fmdzt6Eu9Afp"
      },
      "source": [
        "### Construct the backbone with proper parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "msnxCHSa9Cb2"
      },
      "outputs": [],
      "source": [
        "model_id = 'a0'\n",
        "use_positional_encoding = model_id in {'a3', 'a4', 'a5'}\n",
        "resolution = 172\n",
        "\n",
        "backbone = movinet.Movinet(\n",
        "    model_id=model_id,\n",
        "    causal=True,\n",
        "    conv_type='2plus1d',\n",
        "    se_type='2plus3d',\n",
        "    activation='hard_swish',\n",
        "    gating_activation='hard_sigmoid',\n",
        "    use_positional_encoding=use_positional_encoding,\n",
        "    use_external_states=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iSCnGqY69IzG"
      },
      "source": [
        "### Construct the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n5m9qZ__eo92"
      },
      "outputs": [],
      "source": [
        "# Note: this is a temporary model constructed for the\n",
        "# purpose of loading the pre-trained checkpoint. Only\n",
        "# the backbone will be used to build the custom classifier.\n",
        "\n",
        "model = movinet_model.MovinetClassifier(\n",
        "    backbone,\n",
        "    num_classes=600,\n",
        "    output_states=True)\n",
        "\n",
        "# Create your example input here.\n",
        "# Refer to the paper for recommended input shapes.\n",
        "inputs = tf.ones([1, 13, 172, 172, 3])\n",
        "\n",
        "# [Optional] Build the model and load a pretrained checkpoint.\n",
        "model.build(inputs.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6jfwRvD59RyX"
      },
      "source": [
        "### Load the pretrained weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g7y8zwcz9SJA"
      },
      "outputs": [],
      "source": [
        "# Extract pretrained weights\n",
        "!wget https://storage.googleapis.com/tf_model_garden/vision/movinet/movinet_a0_stream.tar.gz -O movinet_a0_stream.tar.gz -q\n",
        "!tar -xvf movinet_a0_stream.tar.gz\n",
        "\n",
        "checkpoint_dir = 'movinet_a0_stream'\n",
        "checkpoint_path = tf.train.latest_checkpoint(checkpoint_dir)\n",
        "checkpoint = tf.train.Checkpoint(model=model)\n",
        "status = checkpoint.restore(checkpoint_path)\n",
        "status.assert_existing_objects_matched()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X0MyBsZnTrnh"
      },
      "source": [
        "### Set up the distribution strategy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IHyqT0csQvRl"
      },
      "outputs": [],
      "source": [
        "# Detect hardware\n",
        "try:\n",
        "  tpu_resolver = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection\n",
        "except ValueError:\n",
        "  tpu_resolver = None\n",
        "  gpus = tf.config.experimental.list_logical_devices(\"GPU\")\n",
        "\n",
        "# Select appropriate distribution strategy\n",
        "if tpu_resolver:\n",
        "  tf.config.experimental_connect_to_cluster(tpu_resolver)\n",
        "  tf.tpu.experimental.initialize_tpu_system(tpu_resolver)\n",
        "  distribution_strategy = tf.distribute.experimental.TPUStrategy(tpu_resolver)\n",
        "  print('Running on TPU ', tpu_resolver.cluster_spec().as_dict()['worker'])\n",
        "elif len(gpus) \u003e 1:\n",
        "  distribution_strategy = tf.distribute.MirroredStrategy([gpu.name for gpu in gpus])\n",
        "  print('Running on multiple GPUs ', [gpu.name for gpu in gpus])\n",
        "elif len(gpus) == 1:\n",
        "  distribution_strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n",
        "  print('Running on single GPU ', gpus[0].name)\n",
        "else:\n",
        "  distribution_strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n",
        "  print('Running on CPU')\n",
        "\n",
        "print(\"Number of accelerators: \", distribution_strategy.num_replicas_in_sync)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CmPbIku_9ejS"
      },
      "source": [
        "### Construct custom classifier with required number of classes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lpXf9GFWlE2-"
      },
      "outputs": [],
      "source": [
        "def build_classifier(batch_size, num_frames, resolution, backbone, num_classes):\n",
        "  \"\"\"Builds a classifier on top of a backbone model.\"\"\"\n",
        "  model = movinet_model.MovinetClassifier(\n",
        "      backbone=backbone,\n",
        "      num_classes=num_classes)\n",
        "  model.build([batch_size, num_frames, resolution, resolution, 3])\n",
        "\n",
        "  return model\n",
        "\n",
        "# Construct loss, optimizer and compile the model\n",
        "with distribution_strategy.scope():\n",
        "  model = build_classifier(batch_size, num_frames, resolution, backbone, 10)\n",
        "  loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "  optimizer = tf.keras.optimizers.Adam(learning_rate = 0.001)\n",
        "  model.compile(loss=loss_obj, optimizer=optimizer, metrics=['accuracy'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j3S6pmEO9y9w"
      },
      "source": [
        "## Create a callback for storing the checkpoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jw-TR9LBPuP_"
      },
      "outputs": [],
      "source": [
        "checkpoint_path = \"trained_model/cp.ckpt\"\n",
        "checkpoint_dir = os.path.dirname(checkpoint_path)\n",
        "\n",
        "# Create a callback that saves the model's weights\n",
        "cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,\n",
        "                                                 save_weights_only=True,\n",
        "                                                 verbose=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E99ucruN9-cu"
      },
      "source": [
        "## Train the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9c7aikXAvpwz"
      },
      "outputs": [],
      "source": [
        "results = model.fit(train_ds,\n",
        "                    validation_data=val_ds,\n",
        "                    epochs=2,\n",
        "                    validation_freq=1,\n",
        "                    verbose=1,\n",
        "                    callbacks=[cp_callback])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CSgJ8pq2-CTG"
      },
      "source": [
        "## Evaluate the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xtPKfnsNl8SX"
      },
      "outputs": [],
      "source": [
        "model.evaluate(test_ds)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fGckPjcx-GY7"
      },
      "source": [
        "## Plot the test data confusion matrix"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n4aQvflFmBho"
      },
      "outputs": [],
      "source": [
        "def get_actual_predicted_labels(dataset):\n",
        "  \"\"\"\n",
        "    Create a list of actual ground truth values and the predictions from the model.\n",
        "\n",
        "    Args:\n",
        "      dataset: An iterable data structure, such as a TensorFlow Dataset, with features and labels.\n",
        "\n",
        "    Return:\n",
        "      Ground truth and predicted values for a particular dataset.\n",
        "  \"\"\"\n",
        "  actual = [labels for _, labels in dataset.unbatch()]\n",
        "  predicted = model.predict(dataset)\n",
        "\n",
        "  actual = tf.stack(actual, axis=0)\n",
        "  predicted = tf.concat(predicted, axis=0)\n",
        "  predicted = tf.argmax(predicted, axis=1)\n",
        "\n",
        "  return actual, predicted"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aXYqu7QOmEBZ"
      },
      "outputs": [],
      "source": [
        "def plot_confusion_matrix(actual, predicted, labels, ds_type):\n",
        "  cm = tf.math.confusion_matrix(actual, predicted)\n",
        "  ax = sns.heatmap(cm, annot=True, fmt='g')\n",
        "  sns.set(rc={'figure.figsize':(6, 16)})\n",
        "  sns.set(font_scale=1.4)\n",
        "  ax.set_title('Confusion matrix of action recognition for ' + ds_type)\n",
        "  ax.set_xlabel('Predicted Action')\n",
        "  ax.set_ylabel('Actual Action')\n",
        "  plt.xticks(rotation=90)\n",
        "  plt.yticks(rotation=0)\n",
        "  ax.xaxis.set_ticklabels(labels)\n",
        "  ax.yaxis.set_ticklabels(labels)\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SXk55UYDmGxl"
      },
      "outputs": [],
      "source": [
        "fg = FrameGenerator(subset_paths['train'], num_frames, training = True)\n",
        "label_names = list(fg.class_ids_for_name.keys())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QyxeRFB2mKIz"
      },
      "outputs": [],
      "source": [
        "actual, predicted = get_actual_predicted_labels(test_ds)\n",
        "plot_confusion_matrix(actual, predicted, label_names, 'test')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0WLYpxFz-P-m"
      },
      "source": [
        "## Reconstruct the whole model with `use_external_states=True` to make the inference using states."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S29Cemf_Py1h"
      },
      "outputs": [],
      "source": [
        "model_id = 'a0'\n",
        "use_positional_encoding = model_id in {'a3', 'a4', 'a5'}\n",
        "resolution = 172\n",
        "\n",
        "# Create backbone and model.\n",
        "backbone = movinet.Movinet(\n",
        "    model_id=model_id,\n",
        "    causal=True,\n",
        "    conv_type='2plus1d',\n",
        "    se_type='2plus3d',\n",
        "    activation='hard_swish',\n",
        "    gating_activation='hard_sigmoid',\n",
        "    use_positional_encoding=use_positional_encoding,\n",
        "    use_external_states=True,\n",
        ")\n",
        "\n",
        "model = movinet_model.MovinetClassifier(\n",
        "    backbone,\n",
        "    num_classes=10,\n",
        "    output_states=True)\n",
        "\n",
        "# Create your example input here.\n",
        "# Refer to the paper for recommended input shapes.\n",
        "inputs = tf.ones([1, 13, 172, 172, 3])\n",
        "\n",
        "# [Optional] Build the model and load a pretrained checkpoint.\n",
        "model.build(inputs.shape)\n",
        "\n",
        "# Load weights from the checkpoint to the rebuilt model\n",
        "checkpoint_dir = 'trained_model'\n",
        "model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KQDBRqnSIZ23"
      },
      "source": [
        "## Inference using external states"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aN7ExIA9JjK_"
      },
      "outputs": [],
      "source": [
        "def get_top_k(probs, k=5, label_map=CLASSES):\n",
        "  \"\"\"Outputs the top k model labels and probabilities on the given video.\"\"\"\n",
        "  top_predictions = tf.argsort(probs, axis=-1, direction='DESCENDING')[:k]\n",
        "  top_labels = tf.gather(label_map, top_predictions, axis=-1)\n",
        "  top_labels = [label.decode('utf8') for label in top_labels.numpy()]\n",
        "  top_probs = tf.gather(probs, top_predictions, axis=-1).numpy()\n",
        "  return tuple(zip(top_labels, top_probs))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gquiUVQGIZJ2"
      },
      "outputs": [],
      "source": [
        "# Create initial states for the stream model\n",
        "init_states_fn = model.init_states\n",
        "init_states = init_states_fn(tf.shape(tf.ones(shape=[1, 1, 172, 172, 3])))\n",
        "\n",
        "all_logits = []\n",
        "\n",
        "# To run on a video, pass in one frame at a time\n",
        "states = init_states\n",
        "for frames, label in test_ds.take(1):\n",
        "  for clip in frames[0]:\n",
        "    # Input shape: [1, 1, 172, 172, 3]\n",
        "    clip = tf.expand_dims(tf.expand_dims(clip, axis=0), axis=0)\n",
        "    logits, states = model.predict({**states, 'image': clip}, verbose=0)\n",
        "    all_logits.append(logits)\n",
        "\n",
        "logits = tf.concat(all_logits, 0)\n",
        "probs = tf.nn.softmax(logits)\n",
        "\n",
        "final_probs = probs[-1]\n",
        "top_k = get_top_k(final_probs)\n",
        "print()\n",
        "for label, prob in top_k:\n",
        "  print(label, prob)\n",
        "\n",
        "frames, label = list(test_ds.take(1))[0]\n",
        "to_gif(frames[0].numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TYxRHeuUINi8"
      },
      "source": [
        "## Export to saved model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aHaiqMXbVxC6"
      },
      "outputs": [],
      "source": [
        "saved_model_dir = 'model'\n",
        "tflite_filename = 'model.tflite'\n",
        "input_shape = [1, 1, 172, 172, 3]\n",
        "\n",
        "# Convert to saved model\n",
        "export_saved_model.export_saved_model(\n",
        "    model=model,\n",
        "    input_shape=input_shape,\n",
        "    export_path=saved_model_dir,\n",
        "    causal=True,\n",
        "    bundle_input_init_states_fn=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ct2fyC00I63a"
      },
      "source": [
        "## Convert to TF Lite"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lHr__VEsVOea"
      },
      "outputs": [],
      "source": [
        "converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)\n",
        "tflite_model = converter.convert()\n",
        "\n",
        "with open(tflite_filename, 'wb') as f:\n",
        "  f.write(tflite_model)\n",
        "\n",
        "# Create the interpreter and signature runner\n",
        "interpreter = tf.lite.Interpreter(model_path=tflite_filename)\n",
        "runner = interpreter.get_signature_runner()\n",
        "\n",
        "init_states = {\n",
        "    name: tf.zeros(x['shape'], dtype=x['dtype'])\n",
        "    for name, x in runner.get_input_details().items()\n",
        "}\n",
        "del init_states['image']\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bmp9vBZ_I-Bw"
      },
      "source": [
        "## Inference using external states on tflite model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-0Exy6QJgep-"
      },
      "outputs": [],
      "source": [
        "# To run on a video, pass in one frame at a time\n",
        "states = init_states\n",
        "for frames, label in test_ds.take(1):\n",
        "  for clip in frames[0]:\n",
        "    # Input shape: [1, 1, 172, 172, 3]\n",
        "    outputs = runner(**states, image=clip)\n",
        "    logits = outputs.pop('logits')[0]\n",
        "    states = outputs\n",
        "\n",
        "probs = tf.nn.softmax(logits)\n",
        "top_k = get_top_k(probs)\n",
        "print()\n",
        "for label, prob in top_k:\n",
        "  print(label, prob)\n",
        "\n",
        "frames, label = list(test_ds.take(1))[0]\n",
        "to_gif(frames[0].numpy())"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1JB5Zzgt1DPwblAy9ScfajQtyIC-W5__f",
          "timestamp": 1679078074504
        }
      ],
      "toc_visible": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
